# -*- coding: utf-8 -*-
# @Time    : 2020/12/19 下午9:24
# @Author  : lilong

import torch
from torch import nn


"""
参考：https://wjrsbu.smartapps.cn/zhihu/article?from_voters_page=true&id=144904949&oauthType=search&_swebfr=1
"""


# 不带参数module
class CenteredLayer(nn.Module):
    def __init__(self, **kwargs):
        super(CenteredLayer, self).__init__(**kwargs)

    def forward(self, x):
        # return x - x.mean()
        return x - 1


# 使用的时候跟自定义层的使用方法一样，例如：
layer = CenteredLayer()
mm = layer(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float))
print("mm:", mm)

# 或者作为更复杂模型的一部分：
net = nn.Sequential(
    nn.Linear(8, 128),
    CenteredLayer()
)


# 带参数module
class MyDense(nn.Module):
    def __init__(self):
        super(MyDense, self).__init__()
        self.params = nn.Parameter(torch.randn(4, 1))

    def forward(self, x):
        x = torch.mm(x, self.params)
        return x


net = MyDense()
print(net())


class MyDense(nn.Module):
    """
    ParameterList和ParameterDict，这两个类可以传入多个Parameter实例，
    然后可以像python中的列表和字典一样去访问或更改ParameterList或ParameterDict中的Parameter
    """
    def __init__(self):
        super(MyDense, self).__init__()
        self.params = nn.ParameterList([nn.Parameter(torch.randn(4, 4)) for i in range(3)])
        self.params.append(nn.Parameter(torch.randn(4, 4)))
        self.params1 = nn.ParameterDict({
                'linear1': nn.Parameter(torch.randn(4, 4))
        })
        self.params1.update({'linear2': nn.Parameter(torch.randn(4, 2))})   # 新增

    def forward(self, x):
        for i in range(len(self.params)):
            x = torch.mm(x, self.params[i])
        return torch.mm(x, self.params1["linear2"])


net = MyDense()
print(net)
